#!/usr/bin/env python
#
# Simulator components for representing Virtual Organization Clusters
#
# Author:   Mike Murphy <mamurph@cs.clemson.edu>
# Revision: 23 September 2009
#
#   Copyright 2009 Clemson University
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
# This material is based upon work supported under a National Science
# Foundation Graduate Research Fellowship. Any opinions, findings,
# conclusions or recommendations expressed in this publication are those of
# the author(s) and do not necessarily reflect the views of the National
# Science Foundation.

'''
Simulator components for representing Virtual Organization Clusters. These
components provided the necessary mechanisms to simulate VOCs with overlay
scheduling capabilities, without imposing any particular VOC allocation
algorithm (a basic algorithm may be found in the watchdog_sim module). The
VOC class itself uses VMPilot jobs (a subclass of job_model.PilotJob) to
lease resources from a targeted grid site. Virtual machines (virtualization
module) are started on these leased resources to provide virtual clusters.

VOCs are intended to be autonomic in design, expanding and shrinking
themselves in response to changes in the private virtual cluster. The
actual algorithm used to achieve this autonomy is implemented outside this
module (e.g. watchdog_sim.Watchdog), allowing simulations of different
policies to be conducted using the same basic mechanisms.
'''

from cluster_sim import Cluster
from core_common import PJ_START, PJ_STOP, SC_SHUTDOWN, generate_name
from job_model import PilotJob
from results import VMPilotResult
from virtualization import VMImage
from vo_format_reader import VOCollection


class VMPilot(PilotJob):
   '''
   Pilot job for leasing physical machines to run virtual machines. Upon
   starting execution on a host, this job will first request the VM image
   file from the local file cache registered as a brokered resource with
   the site job scheduler. If the file cache and file are available, then
   the image file will be transferred from a server (if necessary) and
   made available via an asynchronous callback. Once the file is available,
   the virtual machine instance will be booted by requesting the "vmm"
   capability of the host machine and issuing a start request. The VM then
   boots and joins an overlay scheduler configured in the VM image.

   The simulator is capable of simulating several possible error conditions
   that could arise in the above procedure. It is possible that the file
   cache resource will not be available from the site. It is also possible
   that the host machine receiving the job will not have a "vmm" capability.
   Either of these cases will result in pilot job termination with an
   error condition.
   '''
   def __init__(self, name, img_name, img_src=None, pr_err=0.0, \
                voc_callback=None, **kwds):
      '''
      @param name: Name of the pilot job
      @type name: str
      @param img_name: Name of the VM image to use
      @type img_name: str
      @param img_src: Reference to a FileStore containing the VM image
      @param pr_err: Probability of a random error in the pilot job (0.0-1.0)
      @type pr_err: float
      @param voc_callback: Callback function to invoke for VOC management
                           purposes. This function will receive two arguments:
                           a reference to the VM instance created by the pilot
                           job, and a state change code ("SC_" prefixed
                           constants in core_common). This function will be
                           called only after the VM has booted.
      @param kwds: Pilot job extra data. If the extra data contains
                   a keyword argument "classAds", then an additional
                   classAd will be added to the classAds value
                   passed into this method: {'IS_METAL': 'TRUE'}
      '''
      if kwds is None:
         kwds = {}
      if kwds.has_key('classAds'):
         kwds['classAds']['IS_METAL'] = 'TRUE'
      else:
         kwds['classAds'] = {'IS_METAL': 'TRUE'}
      PilotJob.__init__(self, self.signal, name=name, vo='_pilot_', \
                        pr_err=pr_err, **kwds)
      self.img_name = img_name
      self.img_src = img_src
      self.abort = False
      self.vm = None  # Set in run: we don't have a VM until job runs
      self.vmm = None  # Also set in run
   def signal(self, job, time, sim, sched, node, signum):
      '''
      Callback method invoked by job_model.PilotJob to receive notification
      of job start, orderly external termination, and disorderly external
      termination (kill).

      @param job: A reference to this pilot job
      @param time: Current simulated time
      @type time: number
      @param sim: A reference to the simulator object
      @param sched: Scheduler used to start this job
      @param node: Worker node on which this job is running
      @param signum: Signal number to be received by the pilot job. See the
                     L{core_common} constants with the "PJ_" prefix, along
                     with the L{job_model.PilotJob} documentation.
      @type signum: int

      @return: a simulator results collection
      '''
      r = []
      if signum == PJ_START:
         r.append(self.run(time, sim, sched))
      else:
         r.append(self.stop(signum))
      return (r,)
   def run(self, time, sim, sched):
      '''
      Initiates the FileCache request for the VM image file, so that the
      boot process may proceed.

      @param time: Current simulated time
      @type time: number
      @param sim: Reference to the simulator object
      @param sched: Reference to the scheduler starting this job

      @rtype: L{results.VMPilotResult}
      '''
      msg = ''
      if not self.abort:
         fc = sched.get_resource('vmstore')
         if fc:
            if self.img_src:
	       fc.add_store(self.img_src)
            fc.get_file(self.img_name, self.boot_vm)
	    msg = 'requested_image'
         else:
            self.error = True
	    self.notify_done(time, sim)
	    msg = 'no_vmstore'
      else:
         msg = 'aborted'
	 self.notify_done(time, sim)
      return VMPilotResult(time, self.img_name, self.name, msg)
   def stop(self, signum):
      '''
      Terminates a running VM or short-circuits the boot process, if no VM
      is currently running.

      @param signum: Termination signal number (use the constants with the
                     "PJ_" prefix in L{core_common})
      @type signum: int

      @return: a VMPilotResult if the job is running, or None of the pilot
      job itself has not yet started
      '''
      if self.vm:
         self.vmm.shut_down_vm(self.vm, signum)
      else:
         self.abort = True
      r = None
      if self.running:
         r = VMPilotResult(self.sim.get_timestamp(), self.img_name, \
                           self.name, 'terminating')
      return r
   def boot_vm(self, time, simobj, img_name, image):
      '''
      Callback method invoked by the site file cache (storage.FileCache)
      object when the VM image file is ready. This method boots the VM
      instance.

      @param time: Current simulated time
      @type time: number
      @param simobj: Reference to the simulator object
      @param img_name: VM image name
      @type img_name: str
      @param image: Reference to the VM image file object

      @rtype: L{results.VMPilotResult}
      '''
      err_msg = 'no_vmm'
      if not self.abort:
         self.vmm = self.worker_node.get_capability('vmm')
         if self.vmm is not None:
            self.vm = self.vmm.boot_vm(image, self.name, \
                                       self.vm_state_callback)
            return VMPilotResult(time, img_name, self.name, 'booting')
         self.error = True
      else:
         err_msg = 'aborted'
      #
      # Aborted before boot or no VMM available
      self.notify_done(time, simobj)
      return VMPilotResult(time, img_name, self.name, err_msg)
   def vm_state_callback(self, sched, state):
      '''
      Callback method invoked by the associated VirtualMachine instance
      whenever the machine state changes. If the state change SC_SHUTDOWN
      is received, then this method will invoke the notify_done() method
      of this object.

      @param sched: Scheduler used to run this job
      @param state: New machine state (see constants in the L{core_common}
                    module with the "SC_" prefix)
      @type state: int

      @rtype: L{results.VMPilotResult}
      '''
      if state == SC_SHUTDOWN:
         self.notify_done(self.sim.get_time(), self.sim)
      return VMPilotResult(self.sim.get_timestamp(), self.img_name, \
                           self.name, 'state callback: %d' % state)
   #
#


class VOC(Cluster):
   '''
   Representation of the mechanism behind a Virtual Organization Cluster.
   Policies for determining when to change cluster size and from what systems
   to lease resources for virtualization are implemented externally (see the
   watchdog_sim module). VOCs are bona fide Clusters and have all the
   properties and methods of Clusters, with the caveat that growing a VOC
   requires a scheduler, site, or grid from which the required resources
   may be obtained.
   '''
   def __init__(self, name, simobj, vo, img_name, img_src, ppe=0.0, **kwds):
      '''
      @param name: Name of the VOC
      @type name: str
      @param simobj: Reference to the simulator object
      @param vo: Virtual Organization to which VOC belongs
      @type vo: str
      @param img_name: Name of the VM image to be used to spawn VOC nodes
      @type img_name: str
      @param img_src: FileStore on which VOC image resides
      @param ppe: Probability of a random pilot job error causing a
                 failure to expand the VOC properly (0.0-1.0)
      @type ppe: float
      @param kwds: Extra data for the VOC
      '''
      Cluster.__init__(self, name, simobj, None, **kwds)
      self.vo = vo
      self.img_name = img_name
      self.img_src = img_src
      self.ppe = ppe
      self.pilot_list = []
   #
   def pilot_notify(self, vm, state):
      '''
      Callback method invoked by the VMPilot jobs to notify the cluster that
      VMs have booted or are being shut down, so that the actual machines may
      be added to, or removed from, the cluster

      @param vm: VirtualMachine whose state is changed
      @param state: State code (uses "SC_" prefixed codes in the
                    L{core_common} module)
      @type state: int
      '''
      if (state == SC_FREE) and (vm not in self.machines):
         self.add_machine(vm)
      elif (state == SC_SHUTDOWN) and (vm in self.machines):
         self.remove_machine(vm)
   #
   def grow(self, count, submit_to, **kwds):
      '''
      Expands a VOC by submitting pilot jobs to lease additional resources.
      Unlike regular Clusters, expansion of a VOC is not instantaneous but is
      instead asynchronous. Pilot jobs will be submitted to the scheduler,
      site, or grid referenced in the submit_to parameter, and then a boot
      delay will occur between the time the jobs are submitted and the time
      at which the VMs are usable.

      @param count: Number of VMs to start
      @type count: int
      @param submit_to: Source from which VM hosting resources are to be
                        leased (instance of Scheduler, Site, Grid, or any
                        other object with an add_submit_event method that
                        permits 2 arguments: time and job)
      @param kwds: Pilot job extra data

      @rtype: list
      @return: a list of *pilot jobs* that have been created (this return
      type differs from the base class implementation)
      '''
      r = []
      basename = 'vm.' + self.name + '-'
      for x in xrange(0, count):
         name = generate_name(basename, self.name_index, self.name_size)
         self.name_index += 1
         pilot = VMPilot(name, self.img_name, img_src=self.img_src, \
                         pr_err=self.ppe, voc_callback=self.pilot_notify, \
                         **kwds)
         submit_to.add_submit_event(self.sim.get_time(), pilot)
         self.pilot_list.append(pilot)
         r.append(pilot)
      return r
   #
   def shrink(self, count, kill_busy=True, sel_pilots=[]):
      '''
      Reduces the size of a VOC by stopping VMs, using the VMPilot stop()
      method to initiate an orderly VM shutdown (PJ_STOP signal). If a list
      of pilot jobs is provided (sel_pilots), then those pilot jobs will be
      stopped first. If count is greater than the length of sel_pilots,
      additional pilot jobs will be selected by checking for pilot jobs
      where the virtual machine has not yet booted or where the VM is not
      busy. Finally, if kill_busy is True, busy VMs will be terminated to
      reach the requested count.

      @param count: Number of VOC nodes to terminate
      @type count: int
      @param kill_busy: Flag specifying whether or not busy VMs will be
                        terminated, killing any running jobs
      @type kill_busy: bool
      @param sel_pilots:   Optional list of specific pilot jobs to terminate.
                           All pilot jobs in the list will be terminated,
                           without checking the length of the list against the
                           requested count parameter. Only once these jobs
                           have been terminated will any other pilot jobs be
                           selected for termination.

      @rtype: list
      @return: a list of *pilot jobs* that have been terminated (differs from
      base class return type). Note that the length of this list may differ
      from count if sel_pilots is longer than count, or if kill_busy is False.
      '''
      r = []
      for pilot in sel_pilots:
         pilot.stop(PJ_STOP)
         r.append(pilot)
      #
      total = len(sel_pilots)
      for pilot in self.pilot_list:
         if not pilot.vm or not pilot.vm.claimed:
            if total < count:
               pilot.stop(PJ_STOP)
               total += 1
               r.append(pilot)
      #
      if kill_busy and (total < count):
         victims = self.sim.random.sample(self.pilot_list, (count - total))
         for victim in victims:
            victim.stop(PJ_STOP)
            total += 1
            r.append(victim)
      #
      for pilot in r:
         self.pilot_list.remove(pilot)
      #
      return r
   #
#


class AutoVOC(VOCollection):
   '''
   Replacement VOCollection object that automatically creates a VOC and
   corresponding grid CE for each VO seen in an input trace.
   '''
   def __init__(self, grid, img_src, default_size, default_sched, \
                default_boot, wd, **wd_kwds):
      '''
      @param grid: Grid to which new VOCs will be added
      @param img_src: FileStore that will be used to contain all VM images
      @param default_size: Size (bytes) of the VM image if no specific size is
                           specified for a particular VO
      @type default_size: int
      @param default_sched: Name of the scheduler to use with the VOC if no
                            specific scheduler is specified
      @type default_sched: str
      @param default_boot: Default VM boot time in simulator time units
      @type default_boot: number
      @param wd: Watchdog class to instantiate to manage the VOC
      @param wd_kwds: Keyword arguments to the watchdog class
      '''
      self.grid = grid
      self.img_src = img_src
      self.default_size = default_size
      self.default_sched = default_sched
      self.default_boot = default_boot
      self.wd = wd
      self.wd_kwds = wd_kwds
   def __setitem__(self, vo, record):
      '''
      Adds a VO to the collection. If the VO is new (i.e. not already in the
      collection), then a new VOC, watchdog instance, and CE are created and
      added to the grid. Otherwise, this method simply changes the watchdog
      policy of an existing VOC.

      @param vo: Name of the Virtual Organization
      @type vo: str
      @param record: VO instance to add to the collection
      '''
      if vo in self:
         # Just a watchdog policy change
         record = self[vo]['_simulator']
         wd = record['watchdog']
         voc = record['voc']
         wd.add_voc(voc, record)
      else:
         # Create a new VOC
         size = self.default_size
         if record.has_key('img_size'):
            size = record['img_size']
         img_name = vo + '.img'
         sched_name = self.default_sched
         if record.has_key('sched_name'):
            sched_name = record['sched_name']
         boot_time = self.default_boot
         if record.has_key('boot_time'):
            boot_time = record['boot_time']
         #
         voc = new_voc_ce(self.grid, vo, img_name, self.img_src, sched_name)
         sched = voc.get_scheduler()
         #
         img = VMImage(img_name, size, sched, boot_time)
         self.img_src.add_file(img)
         #
         kwds = self.wd_kwds.copy()
         kwds['w_sched'] = sched
         kwds['name'] = 'watchdog:' + vo
         wd = self.wd(**kwds)
         wd.add_voc(voc, record)
         #
         pvt = {'voc': voc, 'watchdog': wd, 'img': img}
         record['_simulator'] = pvt
         #
         dict.__setitem__(self, vo, record)
#


def job_rewrite_to_vo(site_name, job):
   '''
   Utility function invoked by the JobReader to rewrite job submission
   targets from the specified site name to the VOC CE dedicated to jobs
   from the given VO.

   @param site_name: Original site name
   @type site_name: str
   @param job: Job object

   @return: the new site name to which the job should be submitted, or None
   if the job lacks VO information
   '''
   new_site = None
   if job.vo != '':
      new_site = 'voc_ce:' + job.vo
   return new_site
#

def new_voc_ce(grid, vo, img_name, img_src, sched_name):
   '''
   Creates a new Virtual Organization Cluster with an associated Compute
   Element and adds both objects to the specified grid. The newly created
   VOC will be empty (no virtual machines will be instantiated). The new
   grid site will be named "voc_ce:VO" where VO is the specified Virtual
   Organization.

   @param grid: Grid object to which new VOC and CE will be added
   @param vo: Name of the Virtual Organization
   @type vo: str
   @param img_name: Name of the VM image used to spawn VOC nodes
   @type img_name: str
   @param img_src: Reference to the FileStore where the VM image is located
   @param sched_name: Name of the scheduler to use with the VOC
   @type sched_name: str

   @return: a reference to the newly created VOC
   '''
   ce_name = 'voc_ce:' + vo
   cluster_name = 'voc:' + vo
   voc = VOC(cluster_name, grid.sim, vo, img_name, img_src)
   grid.register_cluster(voc)
   grid.add_ce(ce_name, cluster_name, sched_name)
   return (voc)
#
